{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://mp.weixin.qq.com/s?__biz=MzkyMDMxNDkxOA==&mid=2247484376&idx=1&sn=adbd2ac34823b4115d57fd4110e39b05&chksm=c195f691f6e27f87093b84364482cd6ab016cc30035cd898ac2838bf53fdca5191afb4e4d106&scene=178&cur_album_id=3447886516519747585#rd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from talib import MA_Type\n",
    "import talib\n",
    "from datetime import datetime, timedelta\n",
    "from utils.indicators import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_hist_k_data(code,start_date,end_date,frequency='d')->pd.DataFrame:\n",
    "    \"\"\"\n",
    "    获取历史K线数据\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    import baostock as bs\n",
    "    bs.login()\n",
    "    rs = bs.query_history_k_data_plus(code,\"date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST\",start_date,end_date,frequency=frequency)\n",
    "    data_list = []\n",
    "    while (rs.error_code == '0') & rs.next():\n",
    "        # 获取一条记录，将记录合并在一起\n",
    "        data_list.append(rs.get_row_data())\n",
    "    result = pd.DataFrame(data_list, columns=rs.fields)\n",
    "    result.open = result.open.astype(float)\n",
    "    result.high = result.high.astype(float)\n",
    "    result.low = result.low.astype(float)\n",
    "    result.close = result.close.astype(float)\n",
    "    result.date= pd.to_datetime(result.date)\n",
    "    result.set_index('date',inplace=True)\n",
    "    bs.logout()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "symbol = 'sz.000001'\n",
    "# 获取K线数据\n",
    "end_time = '2024-07-02'\n",
    "start_time = '2005-01-01'\n",
    "data = get_hist_k_data(symbol, start_time, end_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = data[['open', 'high', 'low', 'close', 'volume']]\n",
    "# 计算威廉指标\n",
    "def calculate_wr(data, period=14):\n",
    "    highest_high = data['high'].rolling(window=period).max()\n",
    "    lowest_low = data['low'].rolling(window=period).min()\n",
    "    wr = -100 * (highest_high - data['close']) / (highest_high - lowest_low)\n",
    "    return wr\n",
    "\n",
    "# 设置超买和超卖阈值\n",
    "overbought_threshold = -20\n",
    "oversold_threshold = -80\n",
    "\n",
    "# 生成交易信号\n",
    "df['WR'] = calculate_wr(df)\n",
    "df['Signal'] = 0\n",
    "df['Position'] = 0\n",
    "\n",
    "# 当WR低于超卖阈值时买入\n",
    "df.loc[df['WR'] < oversold_threshold, 'Signal'] = 1\n",
    "# 当WR高于超买阈值时卖出\n",
    "df.loc[df['WR'] > overbought_threshold, 'Signal'] = -1\n",
    "\n",
    "# 生成交易信号\n",
    "df['signal'] = np.where(df['WR'] > df['Signal'], 1, 0)\n",
    "df['position'] = df['signal'].diff()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 计算策略表现\n",
    "df['strategy_return'] = df['position'].shift(1) * (df['close'] - df['close'].shift(1))\n",
    "\n",
    "# 计算策略收益\n",
    "starting_balance = 10000.0  # 假设初始资金为10000美元\n",
    "df['balance'] = (df['strategy_return'].cumsum() + starting_balance)\n",
    "\n",
    "# 绘制价格和策略收益\n",
    "plt.figure(figsize=(14, 7))\n",
    "plt.plot(df.index, df['close'], label='StockA Price')\n",
    "plt.plot(df.index, df['balance'], label='Strategy Balance')\n",
    "plt.legend()\n",
    "plt.title('StockA Price and Strategy Balance')\n",
    "plt.xlabel('Time')\n",
    "plt.ylabel('Value')\n",
    "plt.show()\n",
    "\n",
    "# 打印最终策略收益\n",
    "print(f\"Final Strategy Balance: {df['balance'].iloc[-1]}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
